package com.peng.leetcode.tree;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * DistanceK
 * 863. 二叉树中所有距离为 K 的结点
 * https://leetcode.cn/problems/all-nodes-distance-k-in-binary-tree/
 *
 * @author: lupeng6
 * @create: 2023/2/4 13:04
 */
public class DistanceK {
    Map<TreeNode, TreeNode> parentMap = new HashMap<>();
    List<Integer> ans = new ArrayList<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        if (root == null) {
            return null;
        }
        parseParent(root);
        findAns(target, null, k, 0);
        return ans;
    }

    private void parseParent(TreeNode root) {
        if (root.left != null) {
            parentMap.put(root.left, root);
            parseParent(root.left);
        }
        if (root.right != null) {
            parentMap.put(root.right, root);
            parseParent(root.right);
        }
    }

    void findAns(TreeNode node, TreeNode from, int k, int depth) {
        if (node == null) {
            return;
        }
        if (k == depth) {
            ans.add(node.val);
            return;
        }
        if (node.left != from) {
            findAns(node.left, node, k, depth + 1);
        }
        if (node.right != from) {
            findAns(node.right, node, k, depth + 1);
        }
        if (parentMap.get(node.val) != from) {
            findAns(parentMap.get(node.val), node, k, depth + 1);
        }
    }
}
